import os
import numpy as np
from sklearn.neighbors import KNeighborsClassifier

def readFile(dirName):
    x = []
    y = []
    #获取文件名列表
    filename_list = getFileName_List(dirName)   # ['TestData/0_0.txt', 'TestData/0_1.txt'....]
    for filename in filename_list:
        # 获取y标签
        arr_filename = filename.split("/")
        y.append(eval(arr_filename[1][0]))
        i = 0
        j = 32
        # 文件中32*32个字符装入zero数组
        arr_x = np.zeros(j * j,dtype="int")
        with open(filename,"rb") as fp:
            while True:
                line = fp.readline().decode("utf-8").strip()
                if line:
                    arr_x[i:j] = list(line)
                    i=i+32
                    j=j+32
                else:
                    x.append(arr_x)
                    break
    return np.array(x),y

def getFileName_List(dirName):
    #获得文件名列表
    filenames = os.listdir(dirName)
    filename_list = [os.path.join(dirName+"/", filename) for filename in filenames]
    return filename_list

if __name__ == '__main__':
    #获取训练集与测试集
    x_train,y_tarin = readFile("TrainData")
    x_test, y_test = readFile("TestData")
    #模型训练
    estimate = KNeighborsClassifier(n_neighbors=5)
    estimate.fit(x_train,y_tarin)
    #模型评估
    score = estimate.score(x_test,y_test)
    print("模型得分：\n",score)
    #预测值与真实值对比
    y_pred = estimate.predict(x_test)
    bool = (np.array(y_pred)==np.array(y_test))
    print("预测值与真实值对比：\n",bool)
